from models.nets import CNNMnist, CNNCifar, CNNCifar100, CNNMnistClip, CNNCifarClip, CNNCifar2, FedAvgCNN
from models.nets_adapter import CNNMnistAdapter, CNNCifarAdapter, CNNCifar100Adapter, SimpleCNNAdapter, SimpleCNN
from models.resnet import resnet18
from models.vit import deit_tiny_patch16_224
import torch


def get_model(config):
    """
    加载模型
    :param args:
    :return:
    """
    if config.model == "mnist_cnn":
        global_model = CNNMnist().to(config.device)
    elif config.model == "fmnist_cnn":
        global_model = CNNMnist().to(config.device)
    elif config.model == "cnn":
        if config.dataset == 'cifar10':
            global_model = FedAvgCNN(in_features=3, num_classes=10, dim=1600).to(config.device)
        elif config.dataset == 'cifar100':
            global_model = FedAvgCNN(in_features=3, num_classes=100, dim=10816).to(config.device)
        else:
            global_model = CNNCifar().to(config.device)
    elif config.model == "svhn_cnn":
        global_model = CNNCifar().to(config.device)
    elif config.model == "cifar100_cnn":
        global_model = CNNCifar100().to(config.device)
    elif config.model == "res":
        # global_model = resnet18()
        global_model = resnet18(num_classes=100).to(config.device)

    elif config.model == "vit":
        global_model = deit_tiny_patch16_224(num_classes=1000,
                                             drop_rate=0.,
                                             drop_path_rate=0.1)
        global_model.head = torch.nn.Linear(global_model.head.in_features, 10)
        global_model = global_model.to(config.device)
        # global_model = torch.nn.DataParallel(global_model)
    elif config.model == "mnist_clip":
        global_model = CNNMnistClip(config).to(config.device)
    elif config.model == "cnn_clip":
        global_model = CNNCifarClip(config).to(config.device)
    elif config.model == "CNNCifar2":
        global_model = CNNCifar2(config).to(config.device)

    elif config.model == "mnist_cnn_adapter":
        global_model = CNNMnistAdapter().to(config.device)
    elif config.model == "fmnist_cnn_adapter":
        global_model = CNNMnistAdapter().to(config.device)
    elif config.model == "cnn_adapter":
        # global_model = CNNCifarAdapter().to(config.device)
        output_dim = 10
        channel = 1
        if config.dataset == 'cifar100':
            output_dim = 100
            channel = 3
        elif config.dataset in ['cifar10', 'svhn']:
            channel = 3
        global_model = SimpleCNNAdapter(input_dim=(16 * 5 * 5), hidden_dims=[120, 84],
                                        output_dim=output_dim, channel=channel).to(config.device)
    elif config.model == "simple_cnn":
        output_dim = 10
        channel = 1
        if config.dataset == 'cifar100':
            output_dim = 100
            channel = 3
        elif config.dataset in ['cifar10', 'svhn']:
            channel = 3
        global_model = SimpleCNN(input_dim=(16 * 5 * 5), hidden_dims=[120, 84], output_dim=output_dim).to(config.device)
    elif config.model == "svhn_cnn_adapter":
        global_model = CNNCifarAdapter().to(config.device)
    elif config.model == "cifar100_cnn_adapter":
        global_model = CNNCifar100Adapter().to(config.device)
    return global_model
